"""
白方控制模块
提供白方控制的统一接口，支持强化学习控制和规则AI控制

作者：wjttdbx
版本：3.0 - 重构版本，模块化设计
"""
import sys
sys.path.insert(0, sys.path[0]+"/../")

import numpy as np
from typing import List, Optional, Union

# 导入新的模块
from sh15.agents.base_agent import BaseAgent, AgentAction, AgentObservation
from sh15.agents.action_executor import ActionExecutor, SmartActionExecutor
from sh15.core.units import WhiteBoat, WhiteDrone, BlackBoat


# ===== 主要接口类 =====

class WhiteAgent(BaseAgent):
    """
    白方智能体主接口
    
    统一的白方智能体接口，可以包装不同类型的智能体实现
    """
    
    def __init__(self, agent_type: str = "balanced", use_smart_executor: bool = False):
        """
        初始化白方智能体
        
        Args:
            agent_type: 智能体类型 ("balanced", "aggressive", "defensive")
            use_smart_executor: 是否使用智能执行器
        """
        super().__init__(f"WhiteAgent_{agent_type}")
        
        # 创建底层智能体
        self.base_agent = create_rule_agent(agent_type)
        
        # 创建动作执行器
        if use_smart_executor:
            self.action_executor = SmartActionExecutor()
        else:
            self.action_executor = ActionExecutor()
    
    def get_actions(self, observation: AgentObservation) -> AgentAction:
        """
        获取动作
        返回的动作格式：
        - 无人艇: [速度, 方向, lock_action, drone_action]
        - 无人机: [速度, 方向, mode]
        """
        return self.base_agent.get_actions(observation)
    
    def reset(self):
        """重置智能体"""
        self.base_agent.reset()
        self.reset_count += 1
    
    def execute_actions(self, action: AgentAction, white_boats: List[WhiteBoat], 
                       white_drones: List[WhiteDrone], core_env, current_time: float, dt: float):
        """执行动作"""
        self.action_executor.execute_actions(action, white_boats, white_drones, core_env, current_time, dt)
    
    def set_debug_mode(self, enabled: bool):
        """设置调试模式"""
        self.action_executor.set_debug_mode(enabled)
    
    def get_agent_info(self) -> dict:
        """获取智能体信息"""
        info = self.get_info()
        info.update({
            'base_agent_type': type(self.base_agent).__name__,
            'executor_type': type(self.action_executor).__name__,
            'base_agent_info': self.base_agent.get_info()
        })
        return info


# ===== 向后兼容接口 =====

class RuleBasedWhiteAgent:
    """
    向后兼容的规则智能体接口
    
    保持与原有代码的兼容性
    """
    
    def __init__(self, strategy: str = "balanced"):
        """初始化规则AI"""
        self.strategy = strategy
        self.agent = create_rule_agent(strategy)
        self.executor = ActionExecutor()
        
    def reset(self):
        """重置AI状态"""
        self.agent.reset()
        
    def get_actions(self, white_boats: List[WhiteBoat], white_drones: List[WhiteDrone], 
                   black_boats: List[BlackBoat], current_time: float) -> np.ndarray:
        """
        获取白方动作（向后兼容接口）
        
        Returns:
            np.ndarray: 动作数组，形状为(5, 4)，每艘无人艇4个动作维度
        """
        # 创建观测对象
        observation = AgentObservation(white_boats, white_drones, black_boats, current_time)
        
        # 获取动作
        action = self.agent.get_actions(observation)
        
        # 返回兼容格式
        return action.get_boat_actions()
    
    def get_boat_actions(self, white_boats: List[WhiteBoat], white_drones: List[WhiteDrone], 
                        black_boats: List[BlackBoat], current_time: float) -> np.ndarray:
        """获取白方无人艇动作"""
        return self.get_actions(white_boats, white_drones, black_boats, current_time)
    
    def get_drone_actions(self, white_boats: List[WhiteBoat], white_drones: List[WhiteDrone], 
                         black_boats: List[BlackBoat], current_time: float) -> np.ndarray:
        """获取白方无人机动作（暂时返回空数组）"""
        flying_drones = [d for d in white_drones if d.state == "flying"]
        return np.zeros((len(flying_drones), 3))


# ===== 特化智能体类 =====

class AggressiveWhiteAgent(RuleBasedWhiteAgent):
    """激进型白方AI（向后兼容）"""
    
    def __init__(self):
        super().__init__("aggressive")


class DefensiveWhiteAgent(RuleBasedWhiteAgent):
    """防守型白方AI（向后兼容）"""
    
    def __init__(self):
        super().__init__("defensive")


# ===== 动作执行器接口 =====

class WhiteActionExecutor:
    """
    白方动作执行器（向后兼容接口）
    """
    
    @staticmethod
    def execute_actions(boat_actions: np.ndarray, drone_actions: np.ndarray, 
                       white_boats: List[WhiteBoat], white_drones: List[WhiteDrone],
                       core_env, current_time: float, dt: float):
        """执行白方动作（新版本接口）"""
        LegacyActionExecutor.execute_actions(
            boat_actions, drone_actions, white_boats, white_drones, core_env, current_time, dt
        )
    
    @staticmethod
    def execute_legacy_actions(actions: np.ndarray, white_boats: List[WhiteBoat], 
                              white_drones: List[WhiteDrone], core_env, 
                              current_time: float, dt: float):
        """执行白方动作（旧版本接口）"""
        LegacyActionExecutor.execute_legacy_actions(
            actions, white_boats, white_drones, core_env, current_time, dt
        )


# ===== 工厂函数 =====

def create_white_agent(agent_type: str = "balanced", **kwargs) -> Union[WhiteAgent, RuleBasedWhiteAgent]:
    """
    创建白方智能体
    
    Args:
        agent_type: 智能体类型
            - "balanced": 平衡策略
            - "aggressive": 激进策略  
            - "defensive": 防守策略
            - "rule_based": 基础规则AI（向后兼容）
        **kwargs: 额外参数
            
    Returns:
        WhiteAgent: 白方智能体实例
    """
    if agent_type == "rule_based":
        # 向后兼容模式
        strategy = kwargs.get("strategy", "balanced")
        return RuleBasedWhiteAgent(strategy)
    else:
        # 新的模块化模式
        use_smart_executor = kwargs.get("use_smart_executor", False)
        return WhiteAgent(agent_type, use_smart_executor)


def create_legacy_agent(agent_type: str = "rule_based") -> RuleBasedWhiteAgent:
    """
    创建兼容旧接口的智能体
    
    Args:
        agent_type: 智能体类型
            
    Returns:
        RuleBasedWhiteAgent: 兼容旧接口的智能体实例
    """
    if agent_type == "aggressive":
        return AggressiveWhiteAgent()
    elif agent_type == "defensive":
        return DefensiveWhiteAgent()
    else:
        return RuleBasedWhiteAgent()


# ===== 便捷函数 =====

def create_observation(white_boats: List[WhiteBoat], white_drones: List[WhiteDrone], 
                      black_boats: List[BlackBoat], current_time: float) -> AgentObservation:
    """创建标准观测对象"""
    return AgentObservation(white_boats, white_drones, black_boats, current_time)


def execute_agent_actions(agent: BaseAgent, observation: AgentObservation, 
                       white_boats: List[WhiteBoat], white_drones: List[WhiteDrone],
                       core_env, current_time: float, dt: float):
    """
    便捷函数：执行智能体的完整决策循环
    
    Args:
        agent: 智能体实例
        observation: 观测信息
        white_boats: 白方无人艇列表
        white_drones: 白方无人机列表
        core_env: 核心环境
        current_time: 当前时间
        dt: 时间步长
    """
    # 获取动作
    action = agent.get_actions(observation)
    
    # 执行动作
    if hasattr(agent, 'execute_actions'):
        agent.execute_actions(action, white_boats, white_drones, core_env, current_time, dt)
    else:
        # 使用默认执行器
        executor = ActionExecutor()
        executor.execute_actions(action, white_boats, white_drones, core_env, current_time, dt)


# ===== 导出接口 =====

__all__ = [
    # 主要接口
    'WhiteAgent',
    'RuleBasedWhiteAgent', 
    'AggressiveWhiteAgent',
    'DefensiveWhiteAgent',
    'WhiteActionExecutor',
    
    # 工厂函数
    'create_white_agent',
    'create_legacy_agent',
    
    # 便捷函数
    'create_observation',
    'execute_agent_actions',
    
    # 数据结构
    'AgentAction',
    'AgentObservation'
] 